package resolvers

import (
	"context"
	"time"

	"github.com/graph-gophers/graphql-go"
	"github.com/pkg/errors"
	"github.com/stackrox/rox/central/graphql/resolvers/embeddedobjs"
	"github.com/stackrox/rox/central/graphql/resolvers/loaders"
	"github.com/stackrox/rox/central/metrics"
	v1 "github.com/stackrox/rox/generated/api/v1"
	"github.com/stackrox/rox/generated/storage"
	"github.com/stackrox/rox/pkg/features"
	pkgMetrics "github.com/stackrox/rox/pkg/metrics"
	"github.com/stackrox/rox/pkg/protocompat"
	"github.com/stackrox/rox/pkg/search"
	"github.com/stackrox/rox/pkg/search/scoped"
	"github.com/stackrox/rox/pkg/utils"
)

func init() {
	schema := getBuilder()
	utils.Must(schema.AddType("ImageComponentV2", []string{
		"architecture: String!",
		"fixedBy: String!",
		"id: ID!",
		"imageId: String!",
		"name: String!",
		"operatingSystem: String!",
		"priority: Int!",
		"riskScore: Float!",
		"source: SourceType!",
		"version: String!",
	}))

	utils.Must(
		// NOTE: This list is and should remain alphabetically ordered
		schema.AddType("ImageComponent", []string{
			"deploymentCount(query: String, scopeQuery: String): Int!",
			"deployments(query: String, scopeQuery: String, pagination: Pagination): [Deployment!]!",
			"fixedBy: String!",
			"id: ID!",
			"imageCount(query: String, scopeQuery: String): Int!",
			"images(query: String, scopeQuery: String, pagination: Pagination): [Image!]!",
			"imageVulnerabilityCount(query: String, scopeQuery: String): Int!",
			"imageVulnerabilityCounter(query: String): VulnerabilityCounter!",
			"imageVulnerabilities(query: String, scopeQuery: String, pagination: Pagination): [ImageVulnerability]!",
			"lastScanned: Time",
			"layerIndex: Int",
			"license: License",
			"location(query: String): String!",
			"name: String!",
			"operatingSystem: String!",
			"plottedImageVulnerabilities(query: String): PlottedImageVulnerabilities!",
			"priority: Int!",
			"riskScore: Float!",
			"source: SourceType!",
			"topImageVulnerability: ImageVulnerability",
			"unusedVarSink(query: String): Int",
			"version: String!",

			// deprecated functions
			"fixedIn: String! @deprecated(reason: \"use 'fixedBy'\")",
		}),
		schema.AddQuery("imageComponent(id: ID): ImageComponent"),
		schema.AddQuery("imageComponents(query: String, scopeQuery: String, pagination: Pagination): [ImageComponent!]!"),
		schema.AddQuery("imageComponentCount(query: String): Int!"),
	)
}

// ImageComponentResolver represents a generic resolver of image component fields.
// Values may come from either an embedded component context, or a top level component context.
// NOTE: This list is and should remain alphabetically ordered
type ImageComponentResolver interface {
	DeploymentCount(ctx context.Context, args RawQuery) (int32, error)
	Deployments(ctx context.Context, args PaginatedQuery) ([]*deploymentResolver, error)
	FixedBy(ctx context.Context) string
	Id(ctx context.Context) graphql.ID
	ImageCount(ctx context.Context, args RawQuery) (int32, error)
	Images(ctx context.Context, args PaginatedQuery) ([]ImageResolver, error)
	ImageVulnerabilityCount(ctx context.Context, args RawQuery) (int32, error)
	ImageVulnerabilityCounter(ctx context.Context, args RawQuery) (*VulnerabilityCounterResolver, error)
	ImageVulnerabilities(ctx context.Context, args PaginatedQuery) ([]ImageVulnerabilityResolver, error)
	LastScanned(ctx context.Context) (*graphql.Time, error)
	LayerIndex() (*int32, error)
	License(ctx context.Context) (*licenseResolver, error)
	Location(ctx context.Context, args RawQuery) (string, error)
	Name(ctx context.Context) string
	OperatingSystem(ctx context.Context) string
	PlottedImageVulnerabilities(ctx context.Context, args RawQuery) (*PlottedImageVulnerabilitiesResolver, error)
	Priority(ctx context.Context) int32
	RiskScore(ctx context.Context) float64
	Source(ctx context.Context) string
	TopImageVulnerability(ctx context.Context) (ImageVulnerabilityResolver, error)
	UnusedVarSink(ctx context.Context, args RawQuery) *int32
	Version(ctx context.Context) string

	// deprecated functions

	FixedIn(ctx context.Context) string
}

// ImageComponent returns an image component based on an input id (name:version)
func (resolver *Resolver) ImageComponent(ctx context.Context, args IDQuery) (ImageComponentResolver, error) {
	defer metrics.SetGraphQLOperationDurationTime(time.Now(), pkgMetrics.Root, "ImageComponent")

	// check permissions
	if err := readImages(ctx); err != nil {
		return nil, err
	}

	// get loader
	loader, err := loaders.GetComponentV2Loader(ctx)
	if err != nil {
		return nil, err
	}

	ret, err := loader.FromID(ctx, string(*args.ID))
	if err != nil {
		return nil, err
	}

	// With flattened model, there can be multiple component IDs for a given component name, version and OS.
	// But the component list page on VM 1.0 groups results by name, version and OS. So on component single page,
	// we should also show component details (like top CVSS, priority etc) and
	// related entities (like images, CVEs and deployments) grouped by component name + version + OS.
	query := search.NewQueryBuilder().
		AddExactMatches(search.Component, ret.GetName()).
		AddExactMatches(search.ComponentVersion, ret.GetVersion()).
		AddExactMatches(search.OperatingSystem, ret.GetOperatingSystem()).
		ProtoQuery()
	componentFlatData, err := resolver.ImageComponentFlatView.Get(ctx, query)
	if err != nil {
		return nil, err
	}

	// TODO(ROX-28808): The ticket referenced is the reason we can get here.  FromID will find
	// a component ID excluded by context but the FlatView will not. We should honor the context. This
	// will be cleaned up with 28808.
	if len(componentFlatData) != 1 {
		return nil, errors.New("unable to find component")
	}

	return resolver.wrapImageComponentV2FlatWithContext(ctx, ret, componentFlatData[0], true, err)
}

// ImageComponents returns image components that match the input query.
func (resolver *Resolver) ImageComponents(ctx context.Context, q PaginatedQuery) ([]ImageComponentResolver, error) {
	defer metrics.SetGraphQLOperationDurationTime(time.Now(), pkgMetrics.Root, "ImageComponents")

	// check permissions
	if err := readImages(ctx); err != nil {
		return nil, err
	}

	// cast query
	query, err := q.AsV1QueryOrEmpty()
	if err != nil {
		return nil, err
	}

	// Get the flattened data
	componentFlatData, err := resolver.ImageComponentFlatView.Get(ctx, query)
	if err != nil {
		return nil, err
	}

	componentIDs := make([]string, 0, len(componentFlatData))
	for _, componentFlat := range componentFlatData {
		componentIDs = append(componentIDs, componentFlat.GetComponentIDs()...)
	}

	// get loader
	loader, err := loaders.GetComponentV2Loader(ctx)
	if err != nil {
		return nil, err
	}

	// Get the Components themselves.  This will be denormalized.  So use the IDs to get them, but use
	// the data returned from Component Flat View to keep order and set just 1 instance of a Component
	componentQuery := search.NewQueryBuilder().AddExactMatches(search.ComponentID, componentIDs...).ProtoQuery()
	componentQuery.Pagination = &v1.QueryPagination{
		SortOptions: query.GetPagination().GetSortOptions(),
	}
	comps, err := loader.FromQuery(ctx, componentQuery)

	// Stash a single instance of a Component to aid in normalizing
	foundComponent := make(map[normalizedImageComponent]*storage.ImageComponentV2)
	for _, comp := range comps {
		normalized := normalizedImageComponent{
			name:    comp.GetName(),
			version: comp.GetVersion(),
			os:      comp.GetOperatingSystem(),
		}
		if _, ok := foundComponent[normalized]; !ok {
			foundComponent[normalized] = comp
		}
	}

	// Normalize the Components based on the flat view to keep them in the correct paging and sort order
	normalizedComponents := make([]*storage.ImageComponentV2, 0, len(componentFlatData))
	for _, componentFlat := range componentFlatData {
		normalized := normalizedImageComponent{
			name:    componentFlat.GetComponent(),
			version: componentFlat.GetVersion(),
			os:      componentFlat.GetOperatingSystem(),
		}
		normalizedComponents = append(normalizedComponents, foundComponent[normalized])
	}

	componentResolvers, err := resolver.wrapImageComponentV2sFlatWithContext(ctx, normalizedComponents, componentFlatData, err)
	if err != nil {
		return nil, err
	}

	// cast as return type
	ret := make([]ImageComponentResolver, 0, len(componentResolvers))
	for _, res := range componentResolvers {
		ret = append(ret, res)
	}
	return ret, nil
}

// ImageComponentCount returns count of image components that match the input query
func (resolver *Resolver) ImageComponentCount(ctx context.Context, args RawQuery) (int32, error) {
	defer metrics.SetGraphQLOperationDurationTime(time.Now(), pkgMetrics.Root, "ImageComponentCount")
	// check permissions
	if err := readImages(ctx); err != nil {
		return 0, err
	}

	// cast query
	query, err := args.AsV1QueryOrEmpty()
	if err != nil {
		return 0, err
	}

	componentCount, err := resolver.ImageComponentFlatView.Count(ctx, query)
	return int32(componentCount), err
}

/*
Sub Resolver Functions
*/

func (resolver *imageComponentV2Resolver) DeploymentCount(ctx context.Context, args RawQuery) (int32, error) {
	defer metrics.SetGraphQLOperationDurationTime(time.Now(), pkgMetrics.ImageComponents, "DeploymentCount")
	return resolver.root.DeploymentCount(resolver.imageComponentScopeContext(ctx), args)
}

func (resolver *imageComponentV2Resolver) Deployments(ctx context.Context, args PaginatedQuery) ([]*deploymentResolver, error) {
	defer metrics.SetGraphQLOperationDurationTime(time.Now(), pkgMetrics.ImageComponents, "Deployments")
	return resolver.root.Deployments(resolver.imageComponentScopeContext(ctx), args)
}

func (resolver *imageComponentV2Resolver) ImageCount(ctx context.Context, args RawQuery) (int32, error) {
	defer metrics.SetGraphQLOperationDurationTime(time.Now(), pkgMetrics.ImageComponents, "ImageCount")
	return resolver.root.ImageCount(resolver.imageComponentScopeContext(ctx), args)
}

func (resolver *imageComponentV2Resolver) Images(ctx context.Context, args PaginatedQuery) ([]ImageResolver, error) {
	defer metrics.SetGraphQLOperationDurationTime(time.Now(), pkgMetrics.ImageComponents, "Images")
	return resolver.root.Images(resolver.imageComponentScopeContext(ctx), args)
}

func (resolver *imageComponentV2Resolver) ImageVulnerabilityCount(ctx context.Context, args RawQuery) (int32, error) {
	defer metrics.SetGraphQLOperationDurationTime(time.Now(), pkgMetrics.ImageComponents, "ImageVulnerabilityCount")
	return resolver.root.ImageVulnerabilityCount(resolver.imageComponentScopeContext(ctx), args)
}

func (resolver *imageComponentV2Resolver) ImageVulnerabilityCounter(ctx context.Context, args RawQuery) (*VulnerabilityCounterResolver, error) {
	defer metrics.SetGraphQLOperationDurationTime(time.Now(), pkgMetrics.ImageComponents, "ImageVulnerabilityCounter")
	return resolver.root.ImageVulnerabilityCounter(resolver.imageComponentScopeContext(ctx), args)
}

func (resolver *imageComponentV2Resolver) ImageVulnerabilities(ctx context.Context, args PaginatedQuery) ([]ImageVulnerabilityResolver, error) {
	defer metrics.SetGraphQLOperationDurationTime(time.Now(), pkgMetrics.ImageComponents, "ImageVulnerabilities")
	if resolver.ctx == nil {
		resolver.ctx = ctx
	}

	return resolver.root.ImageVulnerabilities(resolver.imageComponentScopeContext(ctx), args)
}

func (resolver *imageComponentV2Resolver) LastScanned(ctx context.Context) (*graphql.Time, error) {
	defer metrics.SetGraphQLOperationDurationTime(time.Now(), pkgMetrics.ImageComponents, "LastScanned")
	if resolver.ctx == nil {
		resolver.ctx = ctx
	}

	// Short path. Full image is embedded when image scan resolver is called.
	if scanTime := embeddedobjs.LastScannedFromContext(resolver.ctx); scanTime != nil {
		return &graphql.Time{Time: *scanTime}, nil
	}

	if features.FlattenImageData.Enabled() {
		imageLoader, err := loaders.GetImageV2Loader(resolver.ctx)
		if err != nil {
			return nil, err
		}
		q := search.NewQueryBuilder().AddExactMatches(search.ImageID, resolver.data.GetImageIdV2()).ProtoQuery()

		images, err := imageLoader.FromQuery(resolver.ctx, q)
		if err != nil || len(images) == 0 {
			return nil, err
		} else if len(images) > 1 {
			return nil, errors.New("multiple images matched for last scanned image component query")
		}

		return protocompat.ConvertTimestampToGraphqlTimeOrError(images[0].GetScan().GetScanTime())
	}
	imageLoader, err := loaders.GetImageLoader(resolver.ctx)
	if err != nil {
		return nil, err
	}

	q := search.NewQueryBuilder().AddExactMatches(search.ImageSHA, resolver.data.GetImageId()).ProtoQuery()

	images, err := imageLoader.FromQuery(resolver.ctx, q)
	if err != nil || len(images) == 0 {
		return nil, err
	} else if len(images) > 1 {
		return nil, errors.New("multiple images matched for last scanned image component query")
	}

	return protocompat.ConvertTimestampToGraphqlTimeOrError(images[0].GetScan().GetScanTime())
}

// PlottedImageVulnerabilities returns the data required by top risky entity scatter-plot on vuln mgmt dashboard
func (resolver *imageComponentV2Resolver) PlottedImageVulnerabilities(ctx context.Context, args RawQuery) (*PlottedImageVulnerabilitiesResolver, error) {
	defer metrics.SetGraphQLOperationDurationTime(time.Now(), pkgMetrics.ImageComponents, "PlottedImageVulnerabilities")
	return resolver.root.PlottedImageVulnerabilities(resolver.imageComponentScopeContext(ctx), args)
}

func (resolver *imageComponentV2Resolver) TopImageVulnerability(ctx context.Context) (ImageVulnerabilityResolver, error) {
	defer metrics.SetGraphQLOperationDurationTime(time.Now(), pkgMetrics.ImageComponents, "TopImageVulnerability")
	if resolver.ctx == nil {
		resolver.ctx = ctx
	}

	return resolver.root.TopImageVulnerability(resolver.imageComponentScopeContext(ctx), RawQuery{})
}

func (resolver *imageComponentV2Resolver) LayerIndex() (*int32, error) {
	w, ok := resolver.data.GetHasLayerIndex().(*storage.ImageComponentV2_LayerIndex)
	if !ok {
		return nil, nil
	}
	v := w.LayerIndex
	return &v, nil
}

// Location returns the location of the component.
func (resolver *imageComponentV2Resolver) Location(_ context.Context, _ RawQuery) (string, error) {
	defer metrics.SetGraphQLOperationDurationTime(time.Now(), pkgMetrics.ImageComponents, "Location")

	return resolver.data.GetLocation(), nil
}

func (resolver *imageComponentV2Resolver) UnusedVarSink(_ context.Context, _ RawQuery) *int32 {
	return nil
}

// Following are deprecated functions that are retained to allow UI time to migrate away from them

func (resolver *imageComponentV2Resolver) FixedIn(_ context.Context) string {
	return resolver.data.GetFixedBy()
}

func (resolver *imageComponentV2Resolver) License(_ context.Context) (*licenseResolver, error) {
	return nil, nil
}

/*
Utility Functions
*/

func (resolver *imageComponentV2Resolver) imageComponentScopeContext(ctx context.Context) context.Context {
	if ctx == nil {
		err := utils.ShouldErr(errors.New("argument 'ctx' is nil"))
		if err != nil {
			log.Error(err)
		}
	}
	if resolver.ctx == nil {
		resolver.ctx = ctx
	}
	if resolver.flatData != nil && len(resolver.flatData.GetComponentIDs()) > 0 {
		return scoped.Context(resolver.ctx, scoped.Scope{
			Level: v1.SearchCategory_IMAGE_COMPONENTS_V2,
			IDs:   resolver.flatData.GetComponentIDs(),
		})
	}
	return scoped.Context(resolver.ctx, scoped.Scope{
		Level: v1.SearchCategory_IMAGE_COMPONENTS_V2,
		IDs:   []string{resolver.data.GetId()},
	})
}
